import torch

import numpy as np
import pandas as pd
from sklearn.metrics import confusion_matrix as calculate_cm

from config import args
from timer import Timer

zero_mat_true = torch.zeros((args.BATCH, args.class_num)).cuda()
# otherwise, true_mat and pred_mat will point to the same reference
zero_mat_pred = torch.zeros((args.BATCH, args.class_num)).cuda()

one_mat = torch.ones((args.BATCH, 1)).cuda()
cm_mat = torch.zeros((args.class_num + 1, args.class_num + 1)).int().cuda()


def confusion_matrix(target, pred, cm_base, n_class=args.class_num,
                     img_name=None, wrong_imgs: list = None, **kwargs):
    """
    return: confusion matrix for each batch
    """
    # accumulator = kwargs.get('accumulator', None)

    batch_size = target.size()[0]

    if args.use_focal:
        true = target.unsqueeze(1)
    else:
        true = torch.argmax(target, 1, keepdim=True)  # (BATCH, 1)
    pred = torch.argmax(pred, 1, keepdim=True)  # (BATCH, 1)

    z_mat: torch.Tensor
    if batch_size == args.BATCH:
        z_mat_true = zero_mat_true
        z_mat_pred = zero_mat_pred
        o_mat = one_mat
    else:
        z_mat_true = torch.zeros((batch_size, n_class)).cuda()
        z_mat_pred = torch.zeros((batch_size, n_class)).cuda()
        o_mat = torch.ones((batch_size, 1)).cuda()

    true_mat = z_mat_true.zero_().scatter_(1, true, o_mat)
    pred_mat = z_mat_pred.zero_().scatter_(1, pred, o_mat)

    cm = cm_mat
    cm[:-1, :-1] = true_mat.t().mm(pred_mat).int()
    cm[:-1, -1] = cm[:-1, :-1].sum(1)
    cm[-1, :] = cm[:-1, :].sum(0)

    cm_base += cm

    if img_name is not None and wrong_imgs is not None and len(wrong_imgs) < args.num_wrong_imgs:
        diff_num = args.num_wrong_imgs - len(wrong_imgs)
        diff_num = min(diff_num, batch_size)
        idx = torch.ne(true.squeeze(1), pred.squeeze(1)).nonzero()[:diff_num]
        for i in idx.squeeze(1).cpu().numpy():
            wrong_imgs.append(img_name[i])

    return cm


if __name__ == '__main__':
   pass